package com.liuqi.openai.data.document.splitter;

import com.liuqi.openai.data.document.Document;
import com.liuqi.openai.data.document.DocumentSplitter;
import com.liuqi.openai.data.segment.TextSegment;
import java.util.ArrayList;
import java.util.List;

import static com.liuqi.openai.data.document.splitter.HybridDocumentSplitter.createSegment;

/**
 * 主题文档分段
 *
 * 文档以第一个{@link #splitSymbol}符号进行提取出主题, 剩下的作为文本段.
 * 如果文档中不包含{@link #splitSymbol}符号, 则整个文档进行返回
 *
 * {@link #split(Document)}方法返回的{@link TextSegment}的文本段, 不包含主题.
 *
 * 主题存储在{@link TextSegment#metadata()}中, key = {@link Document#SUBJECT}
 *
 * @author liuqi
 * @date 2025/8/11
 **/
public class SubjectDocumentSplitter implements DocumentSplitter {

    private final static String DEFAULT_SPLIT_SYMBOL = "\n";

    /**
     * 主题文档分割符号 默认值 {@link #DEFAULT_SPLIT_SYMBOL}
     */
    private final String splitSymbol;

    public SubjectDocumentSplitter() {
        this(DEFAULT_SPLIT_SYMBOL);
    }

    public SubjectDocumentSplitter(String splitSymbol) {
        this.splitSymbol = splitSymbol;
    }

    @Override
    public List<TextSegment> split(Document document) {
        List<TextSegment> segments = new ArrayList<>(1);
        String text = document.text();
        int index = text.indexOf(splitSymbol);
        if (index == -1) {
            segments.add(createSegment(text, document));
        } else {
            String subject = text.substring(0, index);
            String doc = text.substring(index + 1);
            TextSegment segment = createSegment(doc, document);
            segment.metadata().add(Document.SUBJECT, subject);

            segments.add(segment);
        }
        return segments;
    }

}
